#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Complete PBPK Modeling Simulation Script

This script implements complete pharmacokinetic processes in the body, including:
- Absorption processes (First-order model and ADAM model)
- Distribution processes (minimal PBPK model and full PBPK model)
- Elimination processes (hepatic metabolism, renal excretion, biliary excretion, etc.)
- Individual simulation (1 person)
- Population simulation (100 people)
- Chart generation and result analysis

Author: BlackCat
Version: 2.0.0
"""

import sys
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from pathlib import Path
from typing import Dict, List, Tuple

# Add src path to system path
# sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../src'))
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
print(f"Project root path: {project_root}")
sys.path.insert(0, project_root)

from src.absorption.absorption_module import AbsorptionModel
from src.absorption.adam_model import ADAMAbsorptionModel
from src.distribution.distribution_module import DistributionModel
from src.elimination.elimination_module import EliminationModel
from src.population.population_data import PopulationData

class IntegratedPBPKSimulator:
    """Complete PBPK Simulator Class"""
    
    def __init__(self):
        """Initialize Simulator"""
        self.population_generator = PopulationData()
        
    def run_complete_simulation(self, 
                              absorption_model: str = 'first_order',
                              distribution_model: str = 'minimal',                              
                              n_subjects: int = 1,
                              population_params: List[Dict] = None,
                              use_hsk7653_params: bool = True) -> Dict:
        """
        Run complete PBPK simulation
        
        Parameters:            
            absorption_model: 'first_order' or 'adam'
            distribution_model: 'minimal' or 'full'
            n_subjects: number of subjects (1 or 100)
            population_params: population parameter list
            
        Returns:
            dict: complete simulation results
        """
        print("=" * 80)
        print("Starting Complete PBPK Simulation")
        print("=" * 80)
        
        # Set simulation parameters
        time_points = np.linspace(0, 48, 200)  # 0-48 hours, 200 points
        
        if use_hsk7653_params:
            # Use HSK7653 specific parameters to match full_pbpk_hsk7653.py
            base_params = {
                'dose': 5.0,  # HSK7653 dose: 5 mg
                'fu': 0.1,      # 游离分数
                'BP': 1.0,      # 血液-血浆浓度比
                'Vss': 0.54,    # 分布容积 (L/kg) - HSK7653 specific
                'Vp': 3.0,      # 血浆容积 (L)
                'Vt': 12.0,     # 组织容积 (L)
                'Ve': 12.0,     # 细胞外液容积 (L)
                'CL_systemic': 0.5,  # 系统清除率 (L/h) - HSK7653 specific
                'CL_hepatic': 0.3,   # 肝脏清除率 (L/h)
                'CL_renal': 0.2,     # 肾脏清除率 (L/h)
                'CL_biliary': 0.05,  # 胆汁清除率 (L/h)
                'CL_other': 0.05,    # 其他清除率 (L/h)
                'solubility': 1.0,   # 溶解度 (mg/mL)
                'permeability': 1e-4, # 渗透性 (cm/s)
                'molecular_weight': 300.0,  # 分子量 (g/mol)
                'pKa': 7.4,     # 解离常数
                'logP': 2.0,    # 脂水分配系数
                'particle_size': 100.0,  # 粒径 (μm)
                'density': 1.2,  # 密度 (g/cm³)
                'fa': 0.998,    # 吸收分数 - HSK7653 specific
                'fg': 1.0,      # 肠道代谢分数
                'fh': 1.0,      # 肝脏代谢分数
                'ka': 2.5,      # 吸收速率常数 - HSK7653 specific
                'bioavailability': 0.998,  # 生物利用度 - HSK7653 specific
                'Kp_liver': 2.0,
                'Kp_kidney': 2.5,
                'Kp_brain': 0.5,
                'Kp_lung': 1.5,
                'Kp_heart': 1.8,
                'Kp_muscle': 1.2,
                'Kp_adipose': 0.8,
                'Kp_spleen': 2.0
            }
        else:
            # Default parameters
            base_params = {
                'dose': 100.0,  # 药量 mg
                'fu': 0.1,      # 游离分数
                'BP': 1.0,      # 血液-血浆浓度比
                'Vss': 15.0,    # 分布容积 (L)
                'Vp': 3.0,      # 血浆容积 (L)
                'Vt': 12.0,     # 组织容积 (L)
                'Ve': 12.0,     # 细胞外液容积 (L)
                'CL_hepatic': 0.5,  # 肝脏清除率 (L/h)
                'CL_renal': 0.3,    # 肾脏清除率 (L/h)
                'CL_biliary': 0.1,  # 胆汁清除率 (L/h)
                'CL_other': 0.05,   # 其他清除率 (L/h)
                'solubility': 1.0,  # 溶解度 (mg/mL)
                'permeability': 1e-4,  # 渗透性 (cm/s)
                'molecular_weight': 300.0,  # 分子量 (g/mol)
                'pKa': 7.4,     # 解离常数
                'logP': 2.0,    # 脂水分配系数
                'particle_size': 100.0,  # 粒径 (μm)
                'density': 1.2,  # 密度 (g/cm³)
                'fa': 1.0,      # 吸收分数
                'fg': 1.0,      # 肠道代谢分数
                'fh': 1.0       # 肝脏代谢分数
            }
        
        results = {
            'absorption_model': absorption_model,
            'distribution_model': distribution_model,            
            'n_subjects': n_subjects,
            'subjects': []
        }
        
        # Determine population parameters to use
        if population_params is None:
            if n_subjects == 1:
                population_params = [base_params]
            else:
                population_params = self.population_generator.generate_population_parameters(n_subjects)
        
        # Run complete simulation for each individual
        for subject_id, subject_params in enumerate(population_params):
            print(f"\nSimulating individual {subject_id + 1}/{len(population_params)}...")
            
            # Merge parameters
            adams_params = {k: v for k, v in subject_params.items() if k in ADAMAbsorptionModel.allowed_params}
            params = {**base_params, **adams_params}
            # params = {**base_params, **subject_params}
            
            # 1. Absorption phase
            absorption_result = self._run_absorption(absorption_model, params, time_points)
            
            # 2. Distribution phase
            distribution_result = self._run_distribution(distribution_model, params, time_points, absorption_result)
            
            # 3. Elimination phase
            elimination_result = self._run_elimination(params, time_points, distribution_result)
            
            # Integrate results
            subject_result = {
                'subject_id': subject_id + 1,
                'parameters': params,
                'absorption': absorption_result,
                'distribution': distribution_result,
                'elimination': elimination_result
            }
            
            results['subjects'].append(subject_result)
            
            if n_subjects == 1:
                print(f"  Absorption complete: {absorption_result['bioavailability']:.2f}% bioavailability")
                print(f"  Distribution complete: {distribution_result['distribution_volume']:.2f} L distribution volume")
                print(f"  Elimination complete: {elimination_result['clearance_rates']['total']:.2f} L/h total clearance")
        
        return results
    
    def _run_absorption(self, absorption_model: str, params: Dict, time_points: np.ndarray) -> Dict:
        """Run absorption phase"""
        if absorption_model == 'first_order':
            model = AbsorptionModel(params)
            result = model.calculate_absorption(params, time_points)
            
            return {
                'model': 'first_order',
                'time_points': time_points,
                'cumulative_absorption': result['cumulative_absorption'],
                'absorption_rate': result['absorption_rate'],
                'plasma_concentration': result['plasma_concentration'],
                'bioavailability': result['bioavailability'],
                'tmax': result['tmax'],
                'cmax': result['cmax'],
                'gastric_emptying': np.zeros_like(time_points),  # Default for first_order
                'intestinal_absorption': np.zeros_like(time_points)  # Default for first_order
            }
            
        elif absorption_model == 'adam':
            model = ADAMAbsorptionModel()
            result = model.calculate_adam_absorption(time_points, params)
            
            return {
                'model': 'adam',
                'time_points': time_points,
                'cumulative_absorption': result['cumulative_absorption'],
                'absorption_rate': result['total_absorption_rate'],
                'plasma_concentration': result.get('plasma_concentration', result['cumulative_absorption']),
                'bioavailability': result['bioavailability'],
                # 'tmax': np.argmax(result['cumulative_absorption']) * 0.5 if len(result['cumulative_absorption']) > 0 else 0,
                'tmax': np.argmax(result['cumulative_absorption']) * (time_points[1] - time_points[0]) if len(result['cumulative_absorption']) > 0 else 0,
                'cmax': np.max(result['cumulative_absorption']) if len(result['cumulative_absorption']) > 0 else 0,
                'gastric_emptying': result.get('regional_amounts', {}).get('stomach', np.zeros_like(time_points)),
                'intestinal_absorption': result.get('regional_absorbed', {}).get('duodenum', np.zeros_like(time_points))
            }
        else:
            # Default to first_order if invalid model specified
            print(f"Warning: Invalid absorption model '{absorption_model}', using 'first_order'")
            return self._run_absorption('first_order', params, time_points)
    
    def _run_distribution(self, distribution_model: str, params: Dict, time_points: np.ndarray, 
                         absorption_result: Dict) -> Dict:
        """Run distribution phase"""
        model = DistributionModel(params)
        
        # Use results from absorption phase as input
        absorption_input = {
            'dose': params['dose'],
            'cumulative': absorption_result['cumulative_absorption'],
            'plasma_concentration': absorption_result['plasma_concentration']
        }
        
        if distribution_model == 'minimal':
            # Simplified distribution model
            result = model.calculate_distribution(params, time_points, absorption_input)
            
            # Calculate brain distribution for consistency
            brain_result = model.calculate_brain_distribution(params, time_points, result['plasma_concentration'])
            
            return {
                'model': 'minimal',
                'time_points': time_points,
                'plasma_concentration': result['plasma_concentration'],
                'distribution_volume': result['distribution_volume'],
                'tissue_concentrations': result['tissue_concentrations'],
                'tissue_partition': result['tissue_partition'],
                'brain_concentration': brain_result['concentration'],
                'brain_partition': brain_result['partition_coefficient']
            }
            
        elif distribution_model == 'full':
            # Full PBPK distribution model - enhanced to match full_pbpk_hsk7653.py
            result = model.calculate_distribution(params, time_points, absorption_input)
            
            # Calculate special tissue distribution
            brain_result = model.calculate_brain_distribution(params, time_points, result['plasma_concentration'])
            
            # Enhanced full PBPK with more accurate compartment modeling
            enhanced_tissue_concentrations = self._enhance_full_pbpk_distribution(
                params, time_points, absorption_result, result['tissue_concentrations']
            )
            
            return {
                'model': 'full',
                'time_points': time_points,
                'plasma_concentration': result['plasma_concentration'],
                'distribution_volume': result['distribution_volume'],
                'tissue_concentrations': enhanced_tissue_concentrations,
                'tissue_partition': result['tissue_partition'],
                'brain_concentration': brain_result['concentration'],
                'brain_partition': brain_result['partition_coefficient'],
                'enhanced_model': True  # Flag to indicate enhanced full PBPK
            }
    
    def _enhance_full_pbpk_distribution(self, params: Dict, time_points: np.ndarray, 
                                      absorption_result: Dict, tissue_concentrations: Dict) -> Dict:
        """
        Enhanced full PBPK distribution modeling to match full_pbpk_hsk7653.py
        
        This method implements accurate compartmental modeling using the
        full PBPK approach from full_pbpk_hsk7653.py
        """
        # Use the exact same approach as full_pbpk_hsk7653.py
        dose = params['dose']
        body_weight = 70.0  # kg (standard adult)
        cardiac_output = 5.0 * body_weight**0.75  # L/h
        plasma_volume = 0.045 * body_weight  # L
        systemic_clearance = params.get('CL_systemic', 0.5)  # L/h
        
        # Exact same compartments as full_pbpk_hsk7653.py
        compartments = {
            'gut_lumen': {'volume': 0.5, 'flow': 0, 'kp': 1.0},
            'stomach': {'volume': 0.1, 'flow': 0.25, 'kp': 1.5},
            'duodenum': {'volume': 0.1, 'flow': 0.35, 'kp': 1.8},
            'jejunum': {'volume': 0.2, 'flow': 0.7, 'kp': 2.0},
            'ileum': {'volume': 0.3, 'flow': 0.7, 'kp': 2.2},
            'colon': {'volume': 0.5, 'flow': 0.2, 'kp': 1.0},
            'liver': {'volume': 0.026 * body_weight, 'flow': 0.23 * body_weight, 'kp': params.get('Kp_liver', 2.0)},
            'kidney': {'volume': 0.004 * body_weight, 'flow': 0.31 * body_weight, 'kp': params.get('Kp_kidney', 2.5)},
            'brain': {'volume': 0.02 * body_weight, 'flow': 0.05 * body_weight, 'kp': params.get('Kp_brain', 0.5)},
            'lung': {'volume': 0.007 * body_weight, 'flow': cardiac_output * 1.0, 'kp': params.get('Kp_lung', 1.5)},
            'heart': {'volume': 0.005 * body_weight, 'flow': cardiac_output * 0.046, 'kp': params.get('Kp_heart', 1.8)},
            'muscle': {'volume': 0.40 * body_weight, 'flow': 0.075 * body_weight, 'kp': params.get('Kp_muscle', 1.2)},
            'adipose': {'volume': 0.21 * body_weight, 'flow': 0.026 * body_weight, 'kp': params.get('Kp_adipose', 0.8)},
            'spleen': {'volume': 0.002 * body_weight, 'flow': cardiac_output * 0.023, 'kp': params.get('Kp_spleen', 2.0)},
            'plasma': {'volume': plasma_volume, 'flow': cardiac_output, 'kp': 1.0}
        }
        
        # Get absorption data - use the exact absorption from ADAM model
        absorption_rate = absorption_result.get('total_absorption_rate', np.zeros_like(time_points))
        cumulative_absorption = absorption_result.get('cumulative_absorption', np.zeros_like(time_points))
        dt = time_points[1] - time_points[0] if len(time_points) > 1 else 0.24
        
        # Initialize concentrations (exact same approach as full_pbpk_hsk7653.py)
        concentrations = {}
        amounts = {}
        
        for comp_name, comp_data in compartments.items():
            concentrations[comp_name] = np.zeros(len(time_points))
            amounts[comp_name] = np.zeros(len(time_points))
        
        # Initial conditions
        gut_volume = compartments['gut_lumen']['volume']
        concentrations['gut_lumen'][0] = dose / gut_volume  # mg/L
        
        # Full PBPK simulation (exact equations from full_pbpk_hsk7653.py)
        for i in range(1, len(time_points)):
            # Absorption from gut to plasma
            if i < len(absorption_rate):
                absorption_rate_val = absorption_rate[i]  # mg/h
            else:
                absorption_rate_val = 0
            
            # Update plasma (central compartment)
            plasma_conc = concentrations['plasma'][i-1]
            plasma_volume = compartments['plasma']['volume']
            
            # Input from absorption
            input_from_absorption = absorption_rate_val / plasma_volume
            
            # Elimination (systemic clearance)
            elimination = (systemic_clearance * plasma_conc / plasma_volume)
            
            # Distribution to organs
            distribution_loss = 0
            for comp_name, comp_data in compartments.items():
                if comp_name not in ['gut_lumen', 'plasma']:
                    flow = comp_data['flow']
                    kp = comp_data['kp']
                    organ_conc = concentrations[comp_name][i-1]
                    organ_volume = comp_data['volume']
                    
                    # Mass transfer: flow * (C_plasma/kp - C_organ) / organ_volume
                    transfer_rate = (flow * (plasma_conc/kp - organ_conc)) / plasma_volume
                    distribution_loss += transfer_rate
            
            # Update plasma concentration (exact equation)
            d_plasma = input_from_absorption - elimination - distribution_loss
            concentrations['plasma'][i] = max(0, plasma_conc + d_plasma * dt)
            
            # Update other organs (exact equation)
            for comp_name, comp_data in compartments.items():
                if comp_name not in ['gut_lumen', 'plasma']:
                    flow = comp_data['flow']
                    kp = comp_data['kp']
                    organ_volume = comp_data['volume']
                    organ_conc = concentrations[comp_name][i-1]
                    
                    # Mass transfer
                    transfer_rate = (flow * (plasma_conc/kp - organ_conc)) / organ_volume
                    concentrations[comp_name][i] = max(0, organ_conc + transfer_rate * dt)
            
            # Update gut lumen (depletion only)
            if i < len(absorption_rate):
                gut_depletion = absorption_rate[i] / gut_volume
                concentrations['gut_lumen'][i] = max(0, concentrations['gut_lumen'][i-1] - gut_depletion * dt)
            else:
                concentrations['gut_lumen'][i] = 0
        
        # Return enhanced concentrations
        enhanced_concentrations = {}
        for tissue in ['liver', 'kidney', 'brain', 'lung', 'heart', 'muscle', 'adipose', 'spleen']:
            enhanced_concentrations[tissue] = concentrations[tissue]
        
        # Merge with existing tissue concentrations
        for tissue in tissue_concentrations:
            if tissue not in enhanced_concentrations:
                enhanced_concentrations[tissue] = tissue_concentrations[tissue]
        
        return enhanced_concentrations
    
    def _run_elimination(self, params: Dict, time_points: np.ndarray, 
                        distribution_result: Dict) -> Dict:
        """Run elimination phase - enhanced to match full_pbpk_hsk7653.py"""
        model = EliminationModel(params)
        
        # Use results from distribution phase as input
        elimination_result = model.calculate_elimination(
            params, time_points, distribution_result
        )
        
        # Calculate detailed clearance rates
        hepatic_clearance = model.calculate_hepatic_clearance(params)
        renal_clearance = model.calculate_renal_clearance(params)
        biliary_clearance = model.calculate_biliary_clearance(params)
        
        # Enhanced clearance calculation matching full_pbpk_hsk7653.py
        systemic_clearance = params.get('CL_systemic', 0.5)  # L/h
        
        # Extract brain data
        brain_concentration = distribution_result['brain_concentration']
        brain_partition = distribution_result['brain_partition']
        
        # Calculate key metrics similar to full_pbpk_hsk7653.py
        plasma_conc = elimination_result['plasma_concentration']
        
        # AUC calculation
        auc = np.trapezoid(plasma_conc, time_points)
        
        # Terminal half-life
        terminal_points = plasma_conc[-10:] if len(plasma_conc) >= 10 else plasma_conc
        if np.any(terminal_points > 0):
            log_conc = np.log(terminal_points[terminal_points > 0])
            time_term = time_points[-len(terminal_points):][terminal_points > 0]
            if len(log_conc) > 1:
                slope, _ = np.polyfit(time_term, log_conc, 1)
                terminal_half_life = -0.693 / slope
            else:
                terminal_half_life = time_points[-1] * 2
        else:
            terminal_half_life = time_points[-1] * 2
        
        return {
            'time_points': time_points,
            'plasma_concentration': elimination_result['plasma_concentration'],
            'clearance_rates': elimination_result['clearance_rates'],
            'elimination_half_life': {
                **elimination_result['elimination_half_life'],
                'terminal': terminal_half_life,
                'systemic_clearance': systemic_clearance
            },
            'tissue_elimination': elimination_result['tissue_elimination'],
            'metabolite_formation': elimination_result['metabolite_formation'],
            'hepatic_clearance': hepatic_clearance,
            'renal_clearance': renal_clearance,
            'biliary_clearance': biliary_clearance,
            'brain_concentration': brain_concentration,
            'brain_partition': brain_partition,
            'auc': auc,
            'cmax': np.max(plasma_conc) if len(plasma_conc) > 0 else 0,
            'tmax': time_points[np.argmax(plasma_conc)] if len(plasma_conc) > 0 else 0
        }

def generate_plots(results: Dict, output_dir: str = "../output/integrated_pbpk_plots"):
    """生成综合图表"""
    Path(output_dir).mkdir(exist_ok=True)
    
    n_subjects = results['n_subjects']
    time_points = results['subjects'][0]['absorption']['time_points']
    
    # Individual charts
    if n_subjects == 1:
        subject = results['subjects'][0]
        
        plt.figure(figsize=(20, 15))
        
        # 1. 吸收过程
        plt.subplot(4, 4, 1)
        plt.plot(time_points, subject['absorption']['cumulative_absorption'], 'b-', linewidth=2)
        plt.title('Absorption Process - Cumulative Absorption')
        plt.xlabel('Time (hours)')
        plt.ylabel('Cumulative Absorption (mg)')
        plt.grid(True)
        
        # 2. 血浆浓度-时间曲线（完整过程）
        plt.subplot(4, 4, 2)
        plt.plot(time_points, subject['absorption']['plasma_concentration'], 'r-', 
                label='After absorption', linewidth=2)
        plt.plot(time_points, subject['distribution']['plasma_concentration'], 'g--', 
                label='After distribution', linewidth=2)
        plt.plot(time_points, subject['elimination']['plasma_concentration'], 'b:', 
                label='Elimination phase', linewidth=2)
        plt.title('Complete PK Process - Plasma Concentration')
        plt.xlabel('Time (hours)')
        plt.ylabel('Plasma Concentration (mg/L)')
        plt.legend()
        plt.grid(True)
        
        # 3. 清除率分布
        plt.subplot(4, 4, 3)
        clearance_data = subject['elimination']['clearance_rates']
        labels = list(clearance_data.keys())
        values = list(clearance_data.values())
        plt.bar(labels, values)
        plt.title('Clearance Distribution')
        plt.ylabel('Clearance (L/h)')
        plt.xticks(rotation=45)
        plt.grid(True)
        
        # 4. 组织分布
        plt.subplot(4, 4, 4)
        for tissue, conc in subject['distribution']['tissue_concentrations'].items():
            if tissue in ['liver', 'kidney', 'brain', 'muscle']:
                plt.plot(time_points[:len(conc)], conc, label=tissue)
        plt.title('Tissue Distribution')
        plt.xlabel('Time (hours)')
        plt.ylabel('Tissue Concentration (mg/kg)')
        plt.legend()
        plt.grid(True)
        
        # 5. 代谢产物形成
        plt.subplot(4, 4, 5)
        metabolite_data = subject['elimination']['metabolite_formation']
        plt.plot(time_points, metabolite_data['metabolite_concentration'], 'g-', linewidth=2)
        plt.title('Metabolite Formation')
        plt.xlabel('Time (hours)')
        plt.ylabel('Metabolite Concentration (mg/L)')
        plt.grid(True)
        
        # 6. 消除半衰期
        plt.subplot(4, 4, 6)
        half_life_data = subject['elimination']['elimination_half_life']
        labels = [k for k, v in half_life_data.items() 
                 if isinstance(v, (int, float)) and k != 'elimination_constant']
        values = [half_life_data[k] for k in labels]
        plt.bar(labels, values)
        plt.title('Elimination Half-life')
        plt.ylabel('Half-life (hours)')
        plt.xticks(rotation=45)
        plt.grid(True)
        
        # 7. 3D PK图
        from mpl_toolkits.mplot3d import Axes3D
        ax = plt.subplot(4, 4, 7, projection='3d')
        ax.plot(time_points, 
               subject['absorption']['cumulative_absorption'], 
               subject['distribution']['plasma_concentration'])
        ax.set_xlabel('Time (hours)')
        ax.set_ylabel('Absorption Amount (mg)')
        ax.set_zlabel('Plasma Concentration (mg/L)')
        ax.set_title('3D PK Relationship')
        
        # 8. 组织消除
        plt.subplot(4, 4, 8)
        for tissue, conc in subject['elimination']['tissue_elimination'].items():
            if tissue in ['liver', 'kidney', 'muscle']:
                plt.plot(time_points[:len(conc)], conc, label=tissue)
        plt.title('Tissue Elimination')
        plt.xlabel('Time (hours)')
        plt.ylabel('Tissue Concentration (mg/kg)')
        plt.grid(True)
        plt.legend()
        
        # 9. 药物总量变化
        plt.subplot(4, 4, 9)
        total_drug = subject['absorption']['cumulative_absorption']
        eliminated = np.cumsum(subject['elimination']['clearance_rates']['total'] * 
                              subject['elimination']['plasma_concentration'] * 0.24)
        plt.plot(time_points[:len(total_drug)], total_drug, 'b-', label='Total amount in body')
        plt.plot(time_points[:len(eliminated)], eliminated[:len(time_points)], 'r--', label='Cumulative elimination')
        plt.title('Total Drug Amount Change')
        plt.xlabel('Time (hours)')
        plt.ylabel('Drug Amount (mg)')
        plt.grid(True)
        plt.legend()
        
        # 9. 个体图表 - 新增胃排空调试曲线
        plt.subplot(4, 4, 10)  # 调整为4行3列布局
        plt.plot(time_points, subject['absorption']['gastric_emptying'], 'm-', linewidth=2, label='Gastric Emptying Rate')
        plt.title('Gastric Emptying Rate')
        plt.xlabel('Time (hours)')
        plt.ylabel('Emptying Rate (fraction/hour)')
        plt.grid(True)
        plt.ylim(0, 1)  # 限制y轴范围为0-1
        plt.legend()

        # 10. 个体图表 - 新增肠道吸收动态
        plt.subplot(4, 4, 11)
        for segment in ['duodenum', 'jejunum_1', 'jejunum_2', 'ileum_1', 'ileum_2', 'ileum_3']:
            plt.plot(time_points, subject['absorption']['intestinal_absorption'], 
                    label=segment.replace('_', ' ').title())
        plt.title('Intestinal Absorption by Segment')
        plt.xlabel('Time (hours)')
        plt.ylabel('Absorption Rate (mg/h)')
        plt.grid(True)
        plt.legend(ncol=2, fontsize='small')        

        # 11. 个体图表 - 新增脑部分布浓度
        plt.subplot(4, 4, 12)
        if 'brain_concentration' in subject['elimination']:
            plt.plot(time_points, subject['elimination']['brain_concentration'], 'c-', linewidth=2, label='Brain Concentration')
        plt.title('Brain Concentration')
        plt.xlabel('Time (hours)')
        plt.ylabel('Concentration (mg/kg)')        
        plt.grid(True)

        # 12. 个体图表 - 新增血脑屏障分配系数
        plt.subplot(4, 4, 13)
        if 'brain_partition' in subject['elimination']:
            plt.bar(['Brain'], subject['elimination']['brain_partition'], color='lightblue', label='Blood-Brain Barrier Partition Coefficient')
        plt.title('Blood-Brain Barrier Partition Coefficient')
        plt.ylabel('Partition Coefficient (BBB Permeability)')
        plt.ylim(0, 1)  # 假设分配系数范围0-1
        plt.grid(axis='y')
        
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, 'individual_complete_analysis.png'), dpi=300, 
                   bbox_inches='tight')
        plt.show()
        plt.close()
    
    # Population charts
    else:
        plt.figure(figsize=(15, 10))
        
        # Plasma concentration distribution
        all_plasma = []
        for subject in results['subjects']:
            all_plasma.append(subject['elimination']['plasma_concentration'])
        
        all_plasma = np.array(all_plasma)
        mean_plasma = np.mean(all_plasma, axis=0)
        std_plasma = np.std(all_plasma, axis=0)
        
        plt.subplot(2, 2, 1)
        for plasma in all_plasma:
            plt.plot(time_points, plasma, alpha=0.3, color='blue')
        plt.plot(time_points, mean_plasma, 'r-', linewidth=2, label='Mean')
        plt.fill_between(time_points, 
                        mean_plasma - 1.96*std_plasma,
                        mean_plasma + 1.96*std_plasma,
                        alpha=0.2, color='red', label='95% CI')
        plt.title('Population Plasma Concentration Distribution')
        plt.xlabel('Time (hours)')
        plt.ylabel('Plasma Concentration (mg/L)')
        plt.legend()
        plt.grid(True)
        
        # Clearance distribution
        plt.subplot(2, 2, 2)
        clearances = [s['elimination']['clearance_rates']['total'] for s in results['subjects']]
        plt.hist(clearances, bins=20, alpha=0.7, color='green', edgecolor='black')
        plt.axvline(np.mean(clearances), color='red', linestyle='dashed', linewidth=2)
        plt.title('Population Total Clearance Distribution')
        plt.xlabel('Total Clearance (L/h)')
        plt.ylabel('Frequency')
        plt.grid(True)
        
        # AUC distribution
        plt.subplot(2, 2, 3)
        aucs = []
        for subject in results['subjects']:
            plasma = subject['elimination']['plasma_concentration']
            auc = np.trapezoid(plasma, time_points)
            aucs.append(auc)
        
        plt.hist(aucs, bins=20, alpha=0.7, color='orange', edgecolor='black')
        plt.axvline(np.mean(aucs), color='red', linestyle='dashed', linewidth=2)
        plt.title('Population AUC Distribution')
        plt.xlabel('AUC (mg·h/L)')
        plt.ylabel('Frequency')
        plt.grid(True)
        
        # Half-life distribution
        plt.subplot(2, 2, 4)
        half_lives = [s['elimination']['elimination_half_life']['total'] 
                     for s in results['subjects']]
        plt.hist(half_lives, bins=20, alpha=0.7, color='purple', edgecolor='black')
        plt.axvline(np.mean(half_lives), color='red', linestyle='dashed', linewidth=2)
        plt.title('Population Elimination Half-life Distribution')
        plt.xlabel('Half-life (hours)')
        plt.ylabel('Frequency')
        plt.grid(True)
        
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, 'population_analysis.png'), dpi=300, 
                   bbox_inches='tight')
        plt.show()
        plt.close()

def main():
    """Main function: Run complete PBPK simulation"""
    print("=" * 80)
    print("Complete PBPK Modeling Simulation System - Enhanced Full PBPK Mode")
    print("=" * 80)
    
    # Create simulator
    simulator = IntegratedPBPKSimulator()
    
    # Run individual simulation with HSK7653 parameters
    print("\n1. Running individual simulation (HSK7653 parameters)...")
    individual_results = simulator.run_complete_simulation(
        absorption_model='first_order',
        distribution_model = 'minimal',
        n_subjects=1,
        use_hsk7653_params=True
    )
    
    # Run population simulation with HSK7653 parameters
    print("\n2. Running population simulation (HSK7653 parameters)...")
    population_params = simulator.population_generator.generate_population_parameters(100)
    population_results = simulator.run_complete_simulation(
        absorption_model='first_order',
        distribution_model='minimal',
        n_subjects=100,
        population_params=population_params,
        use_hsk7653_params=True
    )
    
    # Generate plots
    print("\n3. Generating plots...")
    generate_plots(individual_results, "../output/individual_plots")
    generate_plots(population_results, "../output/population_plots")
    
    # Enhanced summary with key metrics
    print("\n" + "=" * 80)
    print("Enhanced Full PBPK Simulation Complete!")
    print("=" * 80)
    
    # Print key metrics for first subject
    if individual_results['subjects']:
        subject = individual_results['subjects'][0]
        elimination_data = subject['elimination']
        print(f"\nKey Metrics for HSK7653 (5 mg dose):")
        print(f"  Cmax: {elimination_data['cmax']:.2f} mg/L")
        print(f"  Tmax: {elimination_data['tmax']:.1f} hours")
        print(f"  AUC: {elimination_data['auc']:.1f} mg·h/L")
        print(f"  Terminal half-life: {elimination_data['elimination_half_life'].get('terminal', 0):.1f} hours")
        print(f"  Systemic clearance: {elimination_data['elimination_half_life'].get('systemic_clearance', 0):.2f} L/h")
    
    print(f"\nResults saved to:")
    print(f"  Individual: /output/individual_plots/")
    print(f"  Population: /output/population_plots/")
    print("=" * 80)

if __name__ == "__main__":
    main()